"""
MultiQueryRetriever 是 LangChain 中一个强大的检索工具，它通过生成多个查询变体来改进检索结果。以下是其主要使用场景：
主要使用场景
1、处理用户查询的模糊性
当用户查询可能有多重解释时，自动生成多个相关查询变体
例如："Python最新特性"可以扩展为"Python 3.11新功能"、"Python最新版本更新内容"等
2、克服词汇不匹配问题
解决用户查询术语与文档术语不一致的问题
例如："AI"可以扩展为"人工智能"、"机器学习"、"深度学习"等
3、多角度信息检索
当需要从不同角度检索信息时，自动生成互补性查询
例如："电动汽车优缺点"可分解为"电动汽车优势"和"电动汽车缺点"
4、跨语言检索增强
生成不同语言版本的查询（如果支持多语言）
"""
from typing import List
from langchain.retrievers import MultiQueryRetriever
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.vectorstores import Chroma
from langchain_core.output_parsers import BaseOutputParser
from langchain_text_splitters import RecursiveCharacterTextSplitter
from models import get_ds_model_client, get_ollama_embeddings_client

class LineListOutputParser(BaseOutputParser):
    def parse(self, text: str) -> List[str]:
        lines = text.strip().split("\n")
        # 过滤掉空行
        return list(filter(None, lines))

# 构建向量检索器
def get_retriever():
    # Load blog post
    loader = WebBaseLoader("https://lilianweng.github.io/posts/2023-06-23-agent/")
    data = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=0)
    documents = text_splitter.split_documents(data)
    embeddings_client = get_ollama_embeddings_client()
    vector_store = Chroma.from_documents(documents=documents, embedding=embeddings_client)
    return vector_store.as_retriever()

llm = get_ds_model_client()
retriever=get_retriever()

question = "What are the approaches to Task Decomposition?"
# 内部有提示词调用大模型，把问题生成多个不同的提示
retrieval_from_llm = MultiQueryRetriever.from_llm(
    retriever=retriever,
    llm=llm
)
# Results
unique_docs = retrieval_from_llm.invoke("What does the course say about regression?")
print(len(unique_docs))
print(unique_docs)